# Autor chenfeng
#!/usr/bin/env Python
# coding=utf-8

from paddlenlp import Taskflow
from pprint import pprint


schema = {'婆媳':'人名'}
# schema = {'人名':'婆媳'}
# schema = {'人名':['婆媳','夫妻']}

# schema = ['时间', '事件', '原因', '地址']

content_1 = '2022年1月21日，长安镇德丰村祝家浜19号，吴有成与妻子吴玲芬因口角问题发生纠纷后殴打吴玲芬，未有明显伤势。民警到现场对其教育批评，并开具家暴告诫书。'
content_2 = '2020年1月3日，长安镇港湾清风苑63幢2单元，徐爱明与其儿媳杨学萍发生争吵，网格员赶往现场进行协调。'
contents = []
contents.append(content_1)
contents.append(content_2)

# 默认模型无法准确识别目的地和出发地
few_ie = Taskflow('information_extraction', schema=schema)
result = few_ie(contents)
pprint('*'*40 + '默认模型' + '*'*40)
pprint(result)

# 根据少量样本微调后，模型可以正确识别目的地和出发地
few_ie = Taskflow('information_extraction', schema=schema, task_path='checkpoint/model_best')
result_fine = few_ie(contents)
pprint('*'*40 + '微调模型' + '*'*40)
pprint(result_fine)


